from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Dict, List

BLOCK_HEADER_RE = re.compile(r"^\[Problem\s+(\d+)\]\s*$")
SEP_LINE = '---'

@dataclass
class ProblemBlock:
    problem_id: int
    code: str

def parse_answer_file(path: str) -> Dict[int, ProblemBlock]:
    with open(path, "r", encoding="utf-8") as f:
        lines = f.read().splitlines()

    blocks: Dict[int, ProblemBlock] = {}
    cur_id = None
    buf: List[str] = []
    for line in lines:
        m = BLOCK_HEADER_RE.match(line.strip())
        if m:
            if cur_id is not None:
                blocks[cur_id] = ProblemBlock(cur_id, "\n".join(buf).strip() + "\n")
                buf = []
            cur_id = int(m.group(1))
            continue
        if line.strip() == SEP_LINE:
            if cur_id is not None:
                blocks[cur_id] = ProblemBlock(cur_id, "\n".join(buf).strip() + "\n")
                buf = []
                cur_id = None
            continue
        if cur_id is not None:
            buf.append(line)

    if cur_id is not None and buf:
        blocks[cur_id] = ProblemBlock(cur_id, "\n".join(buf).strip() + "\n")

    return blocks
